import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from mpl_toolkits.mplot3d import Axes3D
import time
import logging
from logging.handlers import TimedRotatingFileHandler
import os
import sys
import json
# 在文件开头添加以下代码（在所有import之后）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置支持中文的字体
plt.rcParams['axes.unicode_minus'] = False    # 正确显示负号


# ====================== 传感器模拟模块 ======================
class IMUSensor:
    def __init__(self):
        self.noise_std = 0.02  # 角速度噪声标准差(rad/s)

    def get_data(self, true_attitude, wind_speed):
        """模拟IMU数据：姿态角 + 角速度"""
        wind_factor = wind_speed * 0.01
        noise = np.random.normal(0, self.noise_std + wind_factor, 3)
        return true_attitude + noise


class GPSSensor:
    def __init__(self):
        self.update_freq = 10  # Hz
        self.last_update = time.time()
        self.position_error = 0.5

    def get_data(self, true_position, in_canyon=False):
        """模拟GPS数据：峡谷区域误差增大5倍[1](@ref)"""
        current_time = time.time()
        if current_time - self.last_update < 1 / self.update_freq:
            return None

        self.last_update = current_time
        error_scale = 5.0 if in_canyon else 1.0
        noise = np.random.multivariate_normal(
            [0, 0, 0],
            np.diag([self.position_error * error_scale] * 3)
        )
        return true_position + noise


class VisionSensor:
    def __init__(self):
        self.range = 30  # 最大检测距离(m)

    def detect_obstacles(self, drone_pos, obstacles):
        """视觉障碍物检测：返回(障碍位置, 距离)列表[1](@ref)"""
        detected = []
        for obs in obstacles:
            dist = np.linalg.norm(drone_pos[:2] - obs[:2])
            if dist < self.range:
                detected.append((*obs, dist))
        return detected


class AirSpeedSensor:
    def get_data(self, true_speed, wind_vector):
        """空速计数据：受风场影响[4](@ref)"""
        return true_speed + wind_vector * 0.7 + np.random.normal(0, 0.2, 3)


# ====================== 数据融合核心 ======================
class SensorFusion:
    def __init__(self):
        self.ekf_state = np.zeros(9)  # [x, y, z, vx, vy, vz, roll, pitch, yaw]
        self.P = np.eye(9)  # 协方差矩阵
        self.mode = "GPS"  # 融合模式

    def _build_jacobian(self, dt):
        """构建状态转移矩阵的雅可比矩阵[4](@ref)"""
        F = np.eye(9)
        F[0, 3] = dt  # x 与 vx 的关系
        F[1, 4] = dt  # y 与 vy 的关系
        F[2, 5] = dt  # z 与 vz 的关系
        return F

    def update_mode(self, gps_quality, obstacles_detected):
        """动态切换融合模式[1](@ref)"""
        if gps_quality < 4 or obstacles_detected:
            self.mode = "VISION"
            self.Q = np.diag([0.1, 0.1, 0.1, 0.5, 0.5, 0.5, 0.01, 0.01, 0.01])  # 视觉主导噪声模型
        else:
            self.mode = "GPS"
            self.Q = np.diag([0.01, 0.01, 0.01, 0.1, 0.1, 0.1, 0.05, 0.05, 0.05])  # GPS主导噪声模型

    def ekf_update(self, imu, gps, vision, airspeed):
        """扩展卡尔曼滤波融合多源数据[4](@ref)"""
        dt = 0.02
        F = self._build_jacobian(dt)
        self.ekf_state = F @ self.ekf_state
        self.P = F @ self.P @ F.T + self.Q

        # 根据模式选择数据源
        if self.mode == "GPS" and gps is not None:
            H = np.zeros((3, 9))
            H[:3, :3] = np.eye(3)
            R = np.diag([0.5, 0.5, 0.5])  # GPS噪声协方差
            self._measurement_update(gps, H, R)

        elif self.mode == "VISION":
            if vision:
                H_vis = np.zeros((2, 9))
                H_vis[:2, :2] = np.eye(2)
                R_vis = np.diag([0.3, 0.3])  # 视觉定位噪声
                self._measurement_update(vision[:2], H_vis, R_vis)

            if airspeed is not None:
                H_air = np.zeros((1, 9))
                H_air[0, 2] = 1  # 高度更新
                R_air = np.array([[0.2]])
                self._measurement_update(np.array([self.ekf_state[2]]), H_air, R_air)

        return self.ekf_state.copy()

    def _measurement_update(self, z, H, R):
        """通用测量更新方程[4](@ref)"""
        y = z - H @ self.ekf_state
        S = H @ self.P @ H.T + R
        K = self.P @ H.T @ np.linalg.inv(S)
        self.ekf_state += K @ y
        self.P = (np.eye(9) - K @ H) @ self.P


# ====================== 自适应PID控制器 ======================
class AdaptivePIDController:
    def __init__(self):
        self.Kp = 1.2
        self.Ki = 0.1
        self.Kd = 0.8
        self.wind_gain_adjust = [1.0, 1.0, 1.0]  # 风扰增益调整因子[7](@ref)
        self.prev_error = np.zeros(3)
        self.integral = np.zeros(3)

    def adjust_for_wind(self, wind_speed):
        """根据风速动态调整PID参数[7](@ref)"""
        wind_factor = min(wind_speed / 10.0, 2.0)  # 限制最大增益倍数
        self.wind_gain_adjust = [
            1 + 0.3 * wind_factor,  # Kp增益
            1 + 0.1 * wind_factor,  # Ki增益
            1 + 0.5 * wind_factor  # Kd增益
        ]

    def compute_control(self, setpoint, current, dt):
        """带风扰补偿的PID控制[7,8](@ref)"""
        error = setpoint - current
        derivative = (error - self.prev_error) / dt

        # 应用风扰增益
        kp_adj = self.Kp * self.wind_gain_adjust[0]
        ki_adj = self.Ki * self.wind_gain_adjust[1]
        kd_adj = self.Kd * self.wind_gain_adjust[2]

        # 积分限幅
        self.integral = np.clip(self.integral + error * dt, -5.0, 5.0)

        control = (
                kp_adj * error +
                ki_adj * self.integral +
                kd_adj * derivative
        )
        self.prev_error = error
        return np.clip(control, -10.0, 10.0)  # 控制量限幅


# ====================== 避障模块 ======================
class ObstacleAvoidance:
    def __init__(self):
        self.safety_dist = 2.5  # 最小安全距离

    def rrt_star_replan(self, drone_pos, obstacles, target):
        """RRT*路径重规划简化实现[1](@ref)"""
        if not obstacles:
            return target

        nearest_obs = min(obstacles, key=lambda o: np.linalg.norm(drone_pos[:2] - o[:2]))
        obs_dist = np.linalg.norm(drone_pos[:2] - nearest_obs[:2])

        if obs_dist < self.safety_dist:
            avoid_direction = drone_pos[:2] - nearest_obs[:2]
            avoid_direction /= np.linalg.norm(avoid_direction)  # 单位化
            new_waypoint = nearest_obs[:2] + avoid_direction * self.safety_dist * 1.5
            return np.array([new_waypoint[0], new_waypoint[1], target[2]])  # 保持高度
        return target

# ====================== 增强型可视化系统 ======================
class EnhancedDroneVisualizer:
    def __init__(self):
        # 创建多面板布局
        self.fig = plt.figure(figsize=(16, 12))
        canvas = self.fig.canvas
        canvas.required_interactive_framework = 'TkAgg'

        # 使用GridSpec创建复杂布局
        gs = GridSpec(3, 2, figure=self.fig, width_ratios=[1.5, 1], height_ratios=[1, 1, 1])

        # 3D轨迹图
        self.ax_3d = self.fig.add_subplot(gs[:, 0], projection='3d')
        self._init_3d_axes()

        # 传感器数据面板
        self.ax_sensors = self.fig.add_subplot(gs[0, 1])
        self._init_sensor_axes()

        # 误差分析面板
        self.ax_error = self.fig.add_subplot(gs[1, 1])
        self._init_error_axes()

        # 风场及控制面板
        self.ax_wind = self.fig.add_subplot(gs[2, 1])
        self._init_wind_axes()

        plt.tight_layout()

    def _init_3d_axes(self):
        """初始化3D轨迹图"""
        self.ax_3d.set_xlim([0, 100])
        self.ax_3d.set_ylim([0, 100])
        self.ax_3d.set_zlim([0, 150])
        self.ax_3d.set_xlabel('X (m)')
        self.ax_3d.set_ylabel('Y (m)')
        self.ax_3d.set_zlabel('Z (m)')
        self.ax_3d.set_title('城市峡谷无人机轨迹跟踪 [c1][c8]')

        # 轨迹容器
        self.true_traj, = self.ax_3d.plot([], [], [], 'b-', lw=2, label='真实轨迹')
        self.fused_traj, = self.ax_3d.plot([], [], [], 'r--', lw=2, label='融合轨迹')
        self.gps_traj, = self.ax_3d.plot([], [], [], 'y:', lw=1, label='GPS轨迹')
        self.vision_traj, = self.ax_3d.plot([], [], [], 'g-.', lw=1, label='视觉轨迹')

        # 无人机和障碍物标记
        self.drone_marker = self.ax_3d.scatter([], [], [], c='r', marker='o', s=100)
        self.obstacles_marker = self.ax_3d.scatter([], [], [], c='k', marker='x', s=100)
        self.target_marker = self.ax_3d.scatter([], [], [], c='g', marker='*', s=200)

        self.ax_3d.legend(loc='upper right')

    def _init_sensor_axes(self):
        """初始化传感器数据面板"""
        self.ax_sensors.clear()
        self.ax_sensors.set_title('传感器数据流 [c1][c5]')
        self.ax_sensors.set_xlabel('时间 (s)')
        self.ax_sensors.set_ylabel('数值')
        self.ax_sensors.grid(True)

        # 传感器数据曲线
        self.imu_roll, = self.ax_sensors.plot([], [], 'b-', label='IMU横滚角')
        self.imu_pitch, = self.ax_sensors.plot([], [], 'g-', label='IMU俯仰角')
        self.gps_error, = self.ax_sensors.plot([], [], 'r-', label='GPS误差')
        self.vision_obs, = self.ax_sensors.plot([], [], 'm-', label='最近障碍距离')

        self.ax_sensors.legend()

    def _init_error_axes(self):
        """初始化误差分析面板"""
        self.ax_error.clear()
        self.ax_error.set_title('融合误差分析 [c3][c6]')
        self.ax_error.set_xlabel('时间 (s)')
        self.ax_error.set_ylabel('误差 (m)')
        self.ax_error.grid(True)

        # 误差曲线
        self.pos_error_x, = self.ax_error.plot([], [], 'r-', label='X位置误差')
        self.pos_error_y, = self.ax_error.plot([], [], 'g-', label='Y位置误差')
        self.pos_error_z, = self.ax_error.plot([], [], 'b-', label='Z位置误差')
        self.error_threshold = self.ax_error.axhline(y=1.0, color='gray', linestyle='--', label='误差阈值')

        # 误差热力图背景
        self.error_heatmap = self.ax_error.scatter([], [], c=[], s=100, alpha=0.3,
                                                   cmap='RdYlGn_r', vmin=0, vmax=3)

        self.ax_error.legend()

    def _init_wind_axes(self):
        """初始化风场及控制面板"""
        self.ax_wind.clear()
        self.ax_wind.set_title('风场影响与控制响应 [c8]')
        self.ax_wind.set_xlabel('时间 (s)')
        self.ax_wind.set_ylabel('数值')
        self.ax_wind.grid(True)

        # 风场和控制曲线
        self.wind_speed, = self.ax_wind.plot([], [], 'c-', label='风速(m/s)')
        self.p_gain, = self.ax_wind.plot([], [], 'm-', label='P增益')
        self.d_gain, = self.ax_wind.plot([], [], 'y-', label='D增益')

        # 控制输出
        self.control_x, = self.ax_wind.plot([], [], 'r-', label='X控制输出')
        self.control_y, = self.ax_wind.plot([], [], 'g-', label='Y控制输出')

        self.ax_wind.legend(ncol=2)

    def update_plot(self, true_pos, fused_pos, gps_pos, vision_pos,
                    obstacles, target, sensor_data, error_data, wind_data, control_data):
        """更新所有可视化面板"""
        # 更新3D轨迹
        self._update_3d_plot(true_pos, fused_pos, gps_pos, vision_pos, obstacles, target)

        # 更新传感器面板
        self._update_sensor_panel(sensor_data)

        # 更新误差面板
        self._update_error_panel(error_data)

        # 更新风场控制面板
        self._update_wind_panel(wind_data, control_data)

        plt.draw()
        plt.pause(0.02)

    def _update_3d_plot(self, true_pos, fused_pos, gps_pos, vision_pos, obstacles, target):
        """更新3D轨迹图"""
        # 更新轨迹
        self._update_trajectory(self.true_traj, true_pos)
        self._update_trajectory(self.fused_traj, fused_pos)
        self._update_trajectory(self.gps_traj, gps_pos)
        self._update_trajectory(self.vision_traj, vision_pos)

        # 更新无人机位置
        self.drone_marker._offsets3d = ([true_pos[0]], [true_pos[1]], [true_pos[2]])

        # 更新障碍物位置
        if obstacles:
            obs_pos = np.array(obstacles)
            self.obstacles_marker._offsets3d = (obs_pos[:, 0], obs_pos[:, 1], obs_pos[:, 2])

        # 更新目标点
        self.target_marker._offsets3d = ([target[0]], [target[1]], [target[2]])

    def _update_sensor_panel(self, sensor_data):
        """更新传感器数据面板"""
        time, imu_roll, imu_pitch, gps_error, min_obs_dist = sensor_data

        self._update_line(self.imu_roll, time, imu_roll)
        self._update_line(self.imu_pitch, time, imu_pitch)
        self._update_line(self.gps_error, time, gps_error)
        self._update_line(self.vision_obs, time, min_obs_dist)

        # 自动调整Y轴范围
        all_data = np.concatenate([imu_roll, imu_pitch, gps_error, min_obs_dist])
        y_min, y_max = np.min(all_data), np.max(all_data)
        self.ax_sensors.set_ylim(y_min - 0.5, y_max + 0.5)

    def _update_error_panel(self, error_data):
        """更新误差分析面板"""
        time, error_x, error_y, error_z = error_data

        self._update_line(self.pos_error_x, time, error_x)
        self._update_line(self.pos_error_y, time, error_y)
        self._update_line(self.pos_error_z, time, error_z)

        # 更新误差热力图
        error_magnitude = np.sqrt(np.array(error_x) ** 2 + np.array(error_y) ** 2 + np.array(error_z) ** 2)
        self.error_heatmap.set_offsets(np.column_stack([time, np.zeros(len(time))]))
        self.error_heatmap.set_array(np.array(error_magnitude))

        # 添加误差统计
        mean_error = np.mean(error_magnitude)
        max_error = np.max(error_magnitude)
        self.ax_error.set_title(f'融合误差分析 [c3][c6] (平均: {mean_error:.2f}m, 最大: {max_error:.2f}m)')

    # 修改EnhancedDroneVisualizer类中的_update_wind_panel方法
    def _update_wind_panel(self, wind_data, control_data):
        """更新风场控制面板"""
        # 解包时添加时间序列
        time, wind_speed, p_gain, d_gain = wind_data
        control_x, control_y = control_data

        # 更新曲线数据
        self._update_line(self.wind_speed, time, wind_speed)
        self._update_line(self.p_gain, time, p_gain)
        self._update_line(self.d_gain, time, d_gain)
        self._update_line(self.control_x, time, control_x)
        self._update_line(self.control_y, time, control_y)

        # 添加风速指示
        wind_level = "轻度" if wind_speed[-1] < 10 else "中度" if wind_speed[-1] < 15 else "强"
        self.ax_wind.set_title(f'风场影响与控制响应 (当前风速: {wind_speed[-1]:.1f}m/s, {wind_level}风扰)')

    def _update_trajectory(self, line, new_point):
        """向轨迹添加新点"""
        x, y, z = line.get_data_3d()
        x = np.append(x, new_point[0])
        y = np.append(y, new_point[1])
        z = np.append(z, new_point[2])
        line.set_data(x, y)
        line.set_3d_properties(z)

    def _update_line(self, line, x_data, y_data):
        """更新2D曲线数据"""
        line.set_data(x_data, y_data)
        # 自动调整X轴范围
        if len(x_data) > 0:
            self.ax_sensors.set_xlim(min(x_data), max(x_data) + 1)


# ====================== 主仿真循环 ======================
def simulate_drone_scenario():
    # 初始化系统组件
    imu = IMUSensor()
    gps = GPSSensor()
    vision = VisionSensor()
    airspeed = AirSpeedSensor()
    fusion = SensorFusion()
    controller = AdaptivePIDController()
    obstacle_avoid = ObstacleAvoidance()
    visualizer = EnhancedDroneVisualizer()  # 使用增强型可视化器

    # 环境参数
    true_position = np.array([0.0, 0.0, 50.0])  # 初始位置
    true_attitude = np.zeros(3)
    target_waypoint = np.array([80, 90, 40])
    wind_vector = np.array([0, 8, 0])  # 初始风速8m/s
    building_obstacles = [(30, 40, 60), (60, 70, 80), (75, 50, 40)]  # 建筑物位置

    # 数据记录 - 扩展记录更多类型数据
    true_trajectory = []
    fused_trajectory = []
    gps_trajectory = []
    vision_trajectory = []

    # 传感器数据记录
    time_history = []
    imu_roll_history = []
    imu_pitch_history = []
    gps_error_history = []
    min_obs_dist_history = []

    # 误差数据记录
    pos_error_x_history = []
    pos_error_y_history = []
    pos_error_z_history = []

    # 风控数据记录
    wind_speed_history = []
    p_gain_history = []
    d_gain_history = []
    control_x_history = []
    control_y_history = []

    # 仿真主循环
    start_time = time.time()
    for t in range(1000):  # 1000步仿真
        current_time = time.time() - start_time

        # 场景事件：25秒后增加风速，35秒出现动态障碍
        if t == 250: wind_vector = np.array([0, 12, 0])  # 风切变
        if t == 350: building_obstacles.append((65, 75, 40))  # 广告气球

        # 传感器数据获取
        imu_data = imu.get_data(true_attitude, np.linalg.norm(wind_vector))
        gps_data = gps.get_data(true_position, in_canyon=(30 < true_position[0] < 70))
        vision_data = vision.detect_obstacles(true_position, building_obstacles)
        airspeed_data = airspeed.get_data(np.array([0, 5, 0]), wind_vector)  # 假设真实空速

        # 视觉位置估计（简化版）
        vision_pos = true_position.copy()
        if vision_data:
            closest_obs = min(vision_data, key=lambda x: x[3])  # 找到最近障碍物
            # 简单的位置估计：基于障碍物位置和距离
            vision_pos[:2] = closest_obs[:2] + np.array([closest_obs[3] * 0.2, 0])

        # 数据融合（带模式切换）
        fusion.update_mode(gps_quality=4 if gps_data is not None else 2,
                           obstacles_detected=bool(vision_data))
        fused_state = fusion.ekf_update(imu_data, gps_data,
                                        vision_data[0][:3] if vision_data else None,
                                        airspeed_data)

        # 避障路径重规划
        safe_waypoint = obstacle_avoid.rrt_star_replan(
            true_position, building_obstacles, target_waypoint
        )

        # 飞行控制（姿态+位置）
        controller.adjust_for_wind(np.linalg.norm(wind_vector))
        control_signal = controller.compute_control(
            setpoint=safe_waypoint,
            current=true_position,
            dt=0.02
        )

        # 更新真实状态（物理模拟）
        true_position += control_signal * 0.02
        true_attitude = 0.95 * true_attitude + 0.05 * control_signal[:3]

        # 记录数据
        true_trajectory.append(true_position.copy())
        fused_trajectory.append(fused_state[:3].copy())
        if gps_data is not None:
            gps_trajectory.append(gps_data.copy())
        vision_trajectory.append(vision_pos.copy())

        # 记录传感器数据
        time_history.append(current_time)
        imu_roll_history.append(imu_data[0])  # 横滚角
        imu_pitch_history.append(imu_data[1])  # 俯仰角

        # 计算GPS误差
        if gps_data is not None:
            gps_error = np.linalg.norm(gps_data - true_position)
        else:
            gps_error = 0
        gps_error_history.append(gps_error)

        # 计算最近障碍距离
        min_obs_dist = min([obs[3] for obs in vision_data]) if vision_data else 30
        min_obs_dist_history.append(min_obs_dist)

        # 计算位置误差
        pos_error = fused_state[:3] - true_position
        pos_error_x_history.append(pos_error[0])
        pos_error_y_history.append(pos_error[1])
        pos_error_z_history.append(pos_error[2])

        # 记录风控数据
        wind_speed_history.append(np.linalg.norm(wind_vector))
        p_gain_history.append(controller.Kp * controller.wind_gain_adjust[0])
        d_gain_history.append(controller.Kd * controller.wind_gain_adjust[2])
        control_x_history.append(control_signal[0])
        control_y_history.append(control_signal[1])

        # 每5步更新一次可视化
        if t % 5 == 0:
            # 准备传感器数据
            sensor_data = (
                time_history,
                imu_roll_history,
                imu_pitch_history,
                gps_error_history,
                min_obs_dist_history
            )

            # 准备误差数据
            error_data = (
                time_history,
                pos_error_x_history,
                pos_error_y_history,
                pos_error_z_history
            )

            # 准备风控数据（添加时间序列）
            wind_data = (
                time_history,  # 添加时间序列
                wind_speed_history,
                p_gain_history,
                d_gain_history
            )

            control_data = (
                control_x_history,
                control_y_history
            )

            visualizer.update_plot(
                true_position,
                fused_state[:3],
                gps_trajectory[-1] if gps_trajectory else np.zeros(3),
                vision_pos,
                building_obstacles,
                safe_waypoint,
                sensor_data,
                error_data,
                wind_data,  # 修改后的风控数据
                control_data  # 控制数据
            )

    # 保存轨迹数据
    np.savez('drone_trajectories.npz',
             true=np.array(true_trajectory),
             fused=np.array(fused_trajectory),
             gps=np.array(gps_trajectory))


if __name__ == "__main__":
    simulate_drone_scenario()